import os

import torch.nn as nn
from torch import optim
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

import random
from torchvision import datasets, transforms

import numpy as np
import torch
import os
import matplotlib.pyplot as plt
from LeNet5卷神经网络 import LeNet5

#设置随机数种子，保证结果的可复现
def setup_seed(seed):
    np.random.seed(seed)
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    torch.manual_seed(seed)

#设置随机数种子
setup_seed(26)

#加载数据,ToTensor()将数据转换为tensor
train_data = datasets.MNIST(root='./data', train=True,download=True, transform=transforms.ToTensor())
test_data = datasets.MNIST(root='./data',train=False,download=True, transform=transforms.ToTensor())

#数据切分
train_loader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=64, shuffle=True)

model = LeNet5().to('cpu')

#定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

#开始迭代
epoches = 10
for epoch in range(1,1+epoches):
    model.train()
    total_loss = 0
    for i ,(images,labels) in enumerate(train_loader):
        #to('cpu')把数据放到设备上
        images = images.to('cpu')
        labels = labels.to('cpu')

        #前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)


        #反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        #总损失
        total_loss += loss
    avg_loss = total_loss / len(train_loader)
    print(f'epoch {epoch} avg_loss {avg_loss}')

#保存模型
sava_model = torch.save(model.state_dict(), './model.pth')





